package com.sali.agent;

import com.sali.annotation.CommonPromptDef;
import com.sali.annotation.SafePromptDef;
import com.sali.config.PromptPropertis;
import com.sali.pojo.PromptLLMCommonModel;
import com.sali.pojo.PromptLLMSafeModel;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @description:
 * @author: sali
 * @date: 2024/9/1 下午4:30.
 * @version: 1.0
 */
@Component
public class PromptAgentLoader  implements ApplicationContextAware, ApplicationRunner {


    private Map<String, List<PromptLLMCommonModel>> llmCommonPromptMap = new HashMap<>();

    private Map<String, List<PromptLLMSafeModel>> llmsSafePromptMap = new HashMap<>();

    @Autowired
    private PromptPropertis promptPropertis;

    private ApplicationContext applicationContext;
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    /**
     * description: 程序启动后加载
     **/
    @Override
    public void run(ApplicationArguments args) {
        this.resolvePromptDefs(promptPropertis.getPromptLLMCommonModels(), promptPropertis.getPromptLLMSafeModels());
    }

    public void resolvePromptDefs(List<PromptLLMCommonModel> promptLLMCommonModels, List<PromptLLMSafeModel> promptLLMSafeModels) {
        String[] beanNames = applicationContext.getBeanDefinitionNames();
        for (String beanName : beanNames) {
            Object bean = applicationContext.getBean(beanName);
            Method[] methods = bean.getClass().getDeclaredMethods();
            for (Method method : methods) {
                if (method.isAnnotationPresent(CommonPromptDef.class)) {
                    CommonPromptDef annotation = method.getAnnotation(CommonPromptDef.class);
                    String aiModelName = annotation.aiModelName();
                    this.llmCommonPromptMap.put(bean.getClass().getName(), promptLLMCommonModels.stream().filter(promptLLMCommonModel -> promptLLMCommonModel.getAiModelName().equalsIgnoreCase(aiModelName)).collect(Collectors.toList()));
                } else if (method.isAnnotationPresent(SafePromptDef.class)) {
                    SafePromptDef annotation = method.getAnnotation(SafePromptDef.class);
                    String aiModelName = annotation.aiModelName();
                    this.llmsSafePromptMap.put(bean.getClass().getName(), promptLLMSafeModels.stream().filter(promptLLMSafeModel -> promptLLMSafeModel.getAiModelName().equalsIgnoreCase(aiModelName)).collect(Collectors.toList()));
                }
            }
        }
    }

    public String getLlmCommonPromptContent(String beanName, String businessName, String promptLang) {
        return llmCommonPromptMap.get(beanName).stream()
                .filter(llmCommonPrompt ->
                        businessName.equalsIgnoreCase(llmCommonPrompt.getBusinessName()) &&
                        promptLang.equalsIgnoreCase(llmCommonPrompt.getLang()))
                .map(PromptLLMCommonModel::getPromptContent).findFirst().orElse(null);
    }

    public String getLlmSafePromptContent(String beanName, String businessName, String className, String promptLang) {
        return llmsSafePromptMap.get(beanName).stream()
                .filter(llmSafePrompt ->
                        promptLang.equalsIgnoreCase(llmSafePrompt.getLang()) &&
                        businessName.equalsIgnoreCase(llmSafePrompt.getBusinessName()) &&
                        className.equalsIgnoreCase(llmSafePrompt.getClassName()))
                .map(PromptLLMSafeModel::getPromptContent).findFirst().orElse(null);
    }
}
